{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "0518.ipynb",
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/yananma/5_programs_per_day/blob/master/0518.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cG3aGfalecLG",
        "colab_type": "text"
      },
      "source": [
        "## 8.1 命令式和符号式混合编程"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mNTpWaC-fiWP",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "097c82c1-f9b9-4bb2-f7a6-c23e62432211"
      },
      "source": [
        "def add(a, b):\n",
        "    return a + b \n",
        "\n",
        "def fancy_func(a, b, c, d):\n",
        "    e = add(a, b)\n",
        "    f = add(c, d)\n",
        "    g = add(e, f)\n",
        "    return g \n",
        "\n",
        "fancy_func(1, 2, 3, 4)"
      ],
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "10"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 2
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bzi9vmJESTBa",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 246
        },
        "outputId": "1fc6f351-8693-4522-a2fd-ab066471793d"
      },
      "source": [
        "def add_str():\n",
        "    return '''\n",
        "def add(a, b):\n",
        "    return a + b \n",
        "''' \n",
        "\n",
        "def fancy_func_str():\n",
        "    return ''' \n",
        "def fancy_func(a, b, c, d):\n",
        "    e = add(a, b)\n",
        "    f = add(c, d)\n",
        "    g = add(e, f)\n",
        "    return g \n",
        "''' \n",
        "\n",
        "def evoke_str():\n",
        "    return add_str() + fancy_func_str() + ''' \n",
        "print(fancy_func(1, 2, 3, 4))\n",
        "''' \n",
        "\n",
        "prog = evoke_str()\n",
        "print(prog)\n",
        "y = compile(prog, '', 'exec')\n",
        "exec(y)"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "\n",
            "def add(a, b):\n",
            "    return a + b \n",
            " \n",
            "def fancy_func(a, b, c, d):\n",
            "    e = add(a, b)\n",
            "    f = add(c, d)\n",
            "    g = add(e, f)\n",
            "    return g \n",
            " \n",
            "print(fancy_func(1, 2, 3, 4))\n",
            "\n",
            "10\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VowLuCBpWLRv",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install mxnet-cu100 d2lzh"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lppqThEoUzl4",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 70
        },
        "outputId": "56062123-5151-44df-bd61-f6059b2a6bf4"
      },
      "source": [
        "from mxnet import nd, sym \n",
        "from mxnet.gluon import nn \n",
        "import time \n",
        "\n",
        "def get_net():\n",
        "    net = nn.HybridSequential()\n",
        "    net.add(nn.Dense(256, activation='relu'), \n",
        "        nn.Dense(128, activation='relu'), \n",
        "        nn.Dense(2))\n",
        "\n",
        "    net.initialize()\n",
        "    return net \n",
        "\n",
        "x = nd.random.normal(shape=(1, 512))\n",
        "net = get_net()\n",
        "net(x)"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "\n",
              "[[0.08827586 0.0050518 ]]\n",
              "<NDArray 1x2 @cpu(0)>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 5
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dMin8kZFWEJ4",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 70
        },
        "outputId": "921ffe84-2b03-4137-9f06-5b8503e1ad4c"
      },
      "source": [
        "net.hybridize()\n",
        "net(x)"
      ],
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "\n",
              "[[0.08827586 0.0050518 ]]\n",
              "<NDArray 1x2 @cpu(0)>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 6
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EOm2mFfUYIpM",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 52
        },
        "outputId": "c756f383-275a-4af7-cc36-df9d4bea2915"
      },
      "source": [
        "def benchmark(net, x):\n",
        "    start = time.time()\n",
        "    for i in range(10000):\n",
        "        _ = net(x)\n",
        "    nd.waitall()\n",
        "    return time.time() - start \n",
        "\n",
        "net = get_net()\n",
        "print('before: %.4f sec' % (benchmark(net, x)))\n",
        "net.hybridize()\n",
        "print('after: %.4f sec' % (benchmark(net, x)))"
      ],
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "before: 6.2879 sec\n",
            "after: 2.5578 sec\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lgBldIzuZL_a",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "net.export('my_mlp')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "voQ13eB_Zu7C",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "575bdf45-663d-47e0-b99c-e26e4bdda54e"
      },
      "source": [
        "x = sym.var('data')\n",
        "net(x)"
      ],
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<Symbol dense5_fwd>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 9
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NhgDpyIuaJgI",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class HybridNet(nn.HybridBlock):\n",
        "    def __init__(self, **kwargs):\n",
        "        super(HybridNet, self).__init__(**kwargs)\n",
        "        self.hidden = nn.Dense(10)\n",
        "        self.output = nn.Dense(2)\n",
        "\n",
        "    def hybrid_forward(self, F, x):\n",
        "        print('F: ', F)\n",
        "        print('x: ', x)\n",
        "        x = F.relu(self.hidden(x))\n",
        "        print('hidden: ', x)\n",
        "        return self.output(x)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "g7DEPwqKdxQR",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 210
        },
        "outputId": "864723ad-15e7-49e0-da89-55f85cd52722"
      },
      "source": [
        "net = HybridNet()\n",
        "net.initialize()\n",
        "x = nd.random.normal(shape=(1, 4))\n",
        "net(x)"
      ],
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "F:  <module 'mxnet.ndarray' from '/usr/local/lib/python3.6/dist-packages/mxnet/ndarray/__init__.py'>\n",
            "x:  \n",
            "[[-0.12225834  0.5429998  -0.9469352   0.59643304]]\n",
            "<NDArray 1x4 @cpu(0)>\n",
            "hidden:  \n",
            "[[0.11134676 0.04770704 0.05341475 0.         0.08091211 0.\n",
            "  0.         0.04143535 0.         0.        ]]\n",
            "<NDArray 1x10 @cpu(0)>\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "\n",
              "[[0.00370749 0.00134991]]\n",
              "<NDArray 1x2 @cpu(0)>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 11
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VYsebgnNeJoo",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 210
        },
        "outputId": "76af0d16-9e88-47c8-c1b9-f0fb002d19e1"
      },
      "source": [
        "net(x)"
      ],
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "F:  <module 'mxnet.ndarray' from '/usr/local/lib/python3.6/dist-packages/mxnet/ndarray/__init__.py'>\n",
            "x:  \n",
            "[[-0.12225834  0.5429998  -0.9469352   0.59643304]]\n",
            "<NDArray 1x4 @cpu(0)>\n",
            "hidden:  \n",
            "[[0.11134676 0.04770704 0.05341475 0.         0.08091211 0.\n",
            "  0.         0.04143535 0.         0.        ]]\n",
            "<NDArray 1x10 @cpu(0)>\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "\n",
              "[[0.00370749 0.00134991]]\n",
              "<NDArray 1x2 @cpu(0)>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 12
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Pr3ViDfteXIi",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 122
        },
        "outputId": "d6f913f6-49a6-4b2a-a75b-6362a94b455b"
      },
      "source": [
        "net.hybridize()\n",
        "net(x)"
      ],
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "F:  <module 'mxnet.symbol' from '/usr/local/lib/python3.6/dist-packages/mxnet/symbol/__init__.py'>\n",
            "x:  <Symbol data>\n",
            "hidden:  <Symbol hybridnet0_relu0>\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "\n",
              "[[0.00370749 0.00134991]]\n",
              "<NDArray 1x2 @cpu(0)>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 13
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-3ZrOKzsefMV",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 70
        },
        "outputId": "ea372979-ca7d-4e5f-d305-ce78d4dcbf10"
      },
      "source": [
        "net(x)"
      ],
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "\n",
              "[[0.00370749 0.00134991]]\n",
              "<NDArray 1x2 @cpu(0)>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 14
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CG82U6JmeqIa",
        "colab_type": "text"
      },
      "source": [
        "## 8.2 异步计算"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4rs4lEsbfeZM",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from mxnet import autograd, gluon, nd \n",
        "from mxnet.gluon import loss as gloss, nn \n",
        "import os \n",
        "import subprocess \n",
        "import time "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WF1qkgyxekvq",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 70
        },
        "outputId": "aab9b0f9-5c5a-4ccb-a5f2-eb0ab3172ad5"
      },
      "source": [
        "a = nd.ones((1, 2))\n",
        "b = nd.ones((1, 2))\n",
        "c = a * b + 2 \n",
        "c"
      ],
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "\n",
              "[[3. 3.]]\n",
              "<NDArray 1x2 @cpu(0)>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 16
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uX0-3cP_f5Q7",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Benchmark():\n",
        "    def __init__(self, prefix=None):\n",
        "        self.prefix = prefix + ' ' if prefix else ''\n",
        "\n",
        "    def __enter__(self):\n",
        "        self.start = time.time()\n",
        "\n",
        "    def __exit__(self, *args):\n",
        "        print('%stime: %.4f sec' % (self.prefix, time.time() - self.start))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MTvL3H8xhoZ8",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 105
        },
        "outputId": "c86cfdfb-721b-406c-d210-4a63631fa37e"
      },
      "source": [
        "with Benchmark('Workloads are queued.'):\n",
        "    x = nd.random.uniform(shape=(2000, 2000))\n",
        "    y = nd.dot(x, x).sum()\n",
        "\n",
        "with Benchmark('Workloads are finished.'):\n",
        "    print('sum=', y)"
      ],
      "execution_count": 18,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Workloads are queued. time: 0.0013 sec\n",
            "sum= \n",
            "[1.9998138e+09]\n",
            "<NDArray 1 @cpu(0)>\n",
            "Workloads are finished. time: 0.4028 sec\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SPeHRbZBiMRV",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "e6c54c39-2d3c-46c5-a656-881e2ee4a3a9"
      },
      "source": [
        "with Benchmark():\n",
        "    y = nd.dot(x, x)\n",
        "    y.wait_to_read()"
      ],
      "execution_count": 19,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "time: 0.2891 sec\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MXeSKWRf4qva",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "76a0a2f7-f4f8-48d4-f8e1-fcc5cfaf1c7a"
      },
      "source": [
        "with Benchmark():\n",
        "    y = nd.dot(x, x)\n",
        "    z = nd.dot(x, x)\n",
        "    nd.waitall()"
      ],
      "execution_count": 20,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "time: 0.5709 sec\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "u1mmdCKx41tc",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "9a1e704a-60ee-45fc-e67d-68a0f9791410"
      },
      "source": [
        "with Benchmark():\n",
        "    y = nd.dot(x,x)\n",
        "    y.asnumpy()"
      ],
      "execution_count": 21,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "time: 0.2777 sec\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6VEoO-z0663c",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "890c61d2-a2c5-4461-c647-b378233ac93a"
      },
      "source": [
        "with Benchmark():\n",
        "    y = nd.dot(x, x)\n",
        "    y.norm().asscalar()"
      ],
      "execution_count": 22,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "time: 0.3418 sec\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TwZxJ2rp7Dlr",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 52
        },
        "outputId": "b362ae4e-2a5c-44f9-c780-b1b67db1aca8"
      },
      "source": [
        "with Benchmark('syn.'):\n",
        "    for _ in range(1000):\n",
        "        y = x + 1 \n",
        "        y.wait_to_read()\n",
        "\n",
        "with Benchmark('asyn.'):\n",
        "    for _ in range(1000):\n",
        "        y = x + 1 \n",
        "    nd.waitall()"
      ],
      "execution_count": 23,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "syn. time: 5.1978 sec\n",
            "asyn. time: 5.1456 sec\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mHhK3wqh7q3p",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def data_iter():\n",
        "    start = time.time()\n",
        "    num_batches, batch_size = 100, 1024 \n",
        "    for i in range(num_batches):\n",
        "        x = nd.random.normal(shape=(batch_size, 512))\n",
        "        y = nd.ones((batch_size, ))\n",
        "        yield x, y \n",
        "        if (i + 1) % 50 == 0:\n",
        "            print('batch %d, time %f sec' % (i + 1, time.time() - start))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QxmcRM-K9UPM",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "net = nn.Sequential()\n",
        "net.add(nn.Dense(2048, activation='relu'), \n",
        "    nn.Dense(512, activation='relu'), \n",
        "    nn.Dense(1))\n",
        "net.initialize()\n",
        "trainer = gluon.Trainer(net.collect_params(), 'sgd', {'learning_rate': 0.005})\n",
        "loss = gloss.L2Loss()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZpjphiUE90o7",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def get_mem():\n",
        "    res = subprocess.check_output(['ps', 'u', '-p', str(os.getpid())])\n",
        "    return int(str(res).split()[15]) / 1e3 "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rNy6VScT-Kvj",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "for X, y in data_iter():\n",
        "    break \n",
        "loss(y, net(X)).wait_to_read()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "co_AEUKD-Vfn",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 70
        },
        "outputId": "8f80890e-90a7-4cd2-d99a-f50e2e587d92"
      },
      "source": [
        "l_sum, mem = 0, get_mem()\n",
        "for X, y in data_iter():\n",
        "    with autograd.record():\n",
        "        l = loss(y, net(X))\n",
        "    l_sum += l.mean().asscalar()\n",
        "    l.backward()\n",
        "    trainer.step(X.shape[0])\n",
        "nd.waitall()\n",
        "print('increased memory: %f MB' %(get_mem() - mem))"
      ],
      "execution_count": 28,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "batch 50, time 13.615280 sec\n",
            "batch 100, time 27.390618 sec\n",
            "increased memory: 11.188000 MB\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Tve5o77Q-_GE",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 70
        },
        "outputId": "8ca43604-fdff-4417-f8ab-35e252728081"
      },
      "source": [
        "mem = get_mem()\n",
        "for X, y in data_iter():\n",
        "    with autograd.record():\n",
        "        l = loss(y, net(X))\n",
        "    l.backward()\n",
        "    trainer.step(X.shape[0])\n",
        "nd.waitall()\n",
        "print('increased memory: %f MB' % (get_mem() - mem))"
      ],
      "execution_count": 29,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "batch 50, time 0.131212 sec\n",
            "batch 100, time 0.273308 sec\n",
            "increased memory: 0.004000 MB\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CM6lLfxrBEm8",
        "colab_type": "text"
      },
      "source": [
        "## 8.3 并行计算"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "r4EnZBmd_2md",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import d2lzh as d2l \n",
        "import mxnet as mx \n",
        "from mxnet import nd "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XXu56tcFBDC1",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def run(x):\n",
        "    return [nd.dot(x, x) for _ in range(10)]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "upyInmxxBh7N",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "x_cpu = nd.random.uniform(shape=(2000, 2000))\n",
        "x_gpu = nd.random.uniform(shape=(6000, 6000), ctx=mx.gpu(0))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zoav5oojBzDl",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 52
        },
        "outputId": "69a3a404-55ce-4ec2-86d8-d2c244daa62c"
      },
      "source": [
        "run(x_cpu)\n",
        "run(x_gpu)\n",
        "nd.waitall()\n",
        "\n",
        "with d2l.Benchmark('Run on CPU.'):\n",
        "    run(x_cpu)\n",
        "    nd.waitall()\n",
        "\n",
        "with d2l.Benchmark('Run on GPU.'):\n",
        "    run(x_gpu)\n",
        "    nd.waitall()"
      ],
      "execution_count": 33,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Run on CPU. time: 2.5805 sec\n",
            "Run on GPU. time: 1.8481 sec\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ldNvLeGtCM6O",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "62631928-bdf2-4de7-beca-ada8cf98e9c8"
      },
      "source": [
        "with d2l.Benchmark('Run on CPU and GPU in parallel.'):\n",
        "    run(x_cpu)\n",
        "    run(x_gpu)\n",
        "    nd.waitall()"
      ],
      "execution_count": 34,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Run on CPU and GPU in parallel. time: 3.3715 sec\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_g-2OQTCCdZQ",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 52
        },
        "outputId": "f2f751b6-0fc4-4c55-f485-90f73027354b"
      },
      "source": [
        "def copy_to_cpu(x):\n",
        "    return [y.copyto(mx.cpu()) for y in x]\n",
        "\n",
        "with d2l.Benchmark('Run on GPU.'):\n",
        "    y = run(x_gpu)\n",
        "    nd.waitall()\n",
        "\n",
        "with d2l.Benchmark('Then copy to CPU.'):\n",
        "    copy_to_cpu(y)\n",
        "    nd.waitall()"
      ],
      "execution_count": 35,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Run on GPU. time: 1.8620 sec\n",
            "Then copy to CPU. time: 0.4153 sec\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "82pq462SDelf",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "cbb43102-1486-4fe2-94ca-1354c726ab69"
      },
      "source": [
        "with d2l.Benchmark('Run and copy in parallel.'):\n",
        "    y = run(x_gpu)\n",
        "    copy_to_cpu(y)\n",
        "    nd.waitall()"
      ],
      "execution_count": 36,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Run and copy in parallel. time: 1.9614 sec\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wpc7dAgvDwTv",
        "colab_type": "text"
      },
      "source": [
        "## 8.4 多 GPU 计算"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4webf6XUD-1G",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 298
        },
        "outputId": "da45d3c9-a62c-4bbd-8f74-1c708573de51"
      },
      "source": [
        "!nvidia-smi"
      ],
      "execution_count": 37,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Wed Nov 13 07:53:25 2019       \n",
            "+-----------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 430.50       Driver Version: 418.67       CUDA Version: 10.1     |\n",
            "|-------------------------------+----------------------+----------------------+\n",
            "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
            "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
            "|===============================+======================+======================|\n",
            "|   0  Tesla K80           Off  | 00000000:00:04.0 Off |                    0 |\n",
            "| N/A   64C    P0    71W / 149W |   1934MiB / 11441MiB |      0%      Default |\n",
            "+-------------------------------+----------------------+----------------------+\n",
            "                                                                               \n",
            "+-----------------------------------------------------------------------------+\n",
            "| Processes:                                                       GPU Memory |\n",
            "|  GPU       PID   Type   Process name                             Usage      |\n",
            "|=============================================================================|\n",
            "+-----------------------------------------------------------------------------+\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VvXYVdm2EKJN",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import d2lzh as d2l \n",
        "import mxnet as mx \n",
        "from mxnet.gluon import loss as gloss \n",
        "from mxnet import autograd, nd \n",
        "import time "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rXqpJ8nkGRf7",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "scale = 0.01 \n",
        "W1 = nd.random.normal(scale=scale, shape=(20, 1, 3, 3))\n",
        "b1 = nd.zeros(shape=20)\n",
        "W2 = nd.random.normal(scale=scale, shape=(50, 20, 5, 5))\n",
        "b2 = nd.zeros(shape=50)\n",
        "W3 = nd.random.normal(scale=scale, shape=(800, 128))\n",
        "b3 = nd.zeros(shape=128)\n",
        "W4 = nd.random.normal(scale=scale, shape=(128, 10))\n",
        "b4 = nd.zeros(shape=10)\n",
        "params = [W1, b1, W2, b2, W3, b3, W4, b4]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "A5SCTkz-HB84",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def lenet(X, params):\n",
        "    h1_conv = nd.Convolution(data=X, weight=params[0], bias=params[1], kernel=(3, 3), num_filter=20)\n",
        "    h1_activation = nd.relu(h1_conv)\n",
        "    h1 = nd.Pooling(data=h1_activation, pool_type='avg', kernel=(2, 2), stride=(2, 2))\n",
        "    \n",
        "    h2_conv = nd.Convolution(data=h1, weight=params[2], bias=params[3], kernel=(5, 5), num_filter=50)\n",
        "    h2_activation = nd.relu(h2_conv)\n",
        "    h2 = nd.Pooling(data=h2_activation, pool_type='avg', kernel=(2, 2), stride=(2, 2))\n",
        "    h2 = nd.flatten(h2)\n",
        "    h3_linear = nd.dot(h2, params[4]) + params[5]\n",
        "    h3 = nd.relu(h3_linear)\n",
        "    y_hat = nd.dot(h3, params[6]) + params[7]\n",
        "    return y_hat \n",
        "\n",
        "loss = gloss.SoftmaxCrossEntropyLoss()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UHVb1aoHIyWh",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def get_params(params, ctx):\n",
        "    new_params = [p.copyto(ctx) for p in params]\n",
        "    for p in new_params:\n",
        "        p.attach_grad()\n",
        "    return new_params "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "T1OdRHVsJGkI",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 122
        },
        "outputId": "9f61dd67-6d9d-448a-900a-31beb61aca53"
      },
      "source": [
        "new_params = get_params(params, mx.gpu(0))\n",
        "print('b1 weight:', new_params[1])\n",
        "print('b1 grad:', new_params[1].grad)"
      ],
      "execution_count": 42,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "b1 weight: \n",
            "[0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]\n",
            "<NDArray 20 @gpu(0)>\n",
            "b1 grad: \n",
            "[0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]\n",
            "<NDArray 20 @gpu(0)>\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cQqXZ0xfJH6a",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def allreduce(data):\n",
        "    for i in range(1, len(data)):\n",
        "        data[0][:] += data[i].copyto(data[0].context)\n",
        "    for i in range(1, len(data)):\n",
        "        data[0].copyto(data[i])"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KIIbRSsnLKiU",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 122
        },
        "outputId": "9b7e17ca-dd37-4250-b630-7232c46c6a65"
      },
      "source": [
        "data = [nd.ones((1, 2), ctx=mx.gpu(i)) * (i + 1) for i in range(1)]\n",
        "print('before: ', data)\n",
        "allreduce(data)\n",
        "print('after: ', data)"
      ],
      "execution_count": 44,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "before:  [\n",
            "[[1. 1.]]\n",
            "<NDArray 1x2 @gpu(0)>]\n",
            "after:  [\n",
            "[[1. 1.]]\n",
            "<NDArray 1x2 @gpu(0)>]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-_jFk5QzLmF5",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def split_and_load(data, ctx):\n",
        "    n, k = data.shape[0], len(ctx)\n",
        "    m = n // k \n",
        "    assert m * k == n, '# examples is not divided by # devices.'\n",
        "    return [data[i * m: (i + 1) * m].as_in_context(ctx[i]) for i in range(k)]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qsp3Wrnsh2Q7",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 316
        },
        "outputId": "637b39c1-23aa-48a1-e7a9-38a941b77366"
      },
      "source": [
        "batch = nd.arange(24).reshape((6, 4))\n",
        "ctx = [mx.gpu(0)]\n",
        "splitted = split_and_load(batch, ctx)\n",
        "print('input: ', batch)\n",
        "print('load into', ctx)\n",
        "print('output:', splitted)"
      ],
      "execution_count": 46,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "input:  \n",
            "[[ 0.  1.  2.  3.]\n",
            " [ 4.  5.  6.  7.]\n",
            " [ 8.  9. 10. 11.]\n",
            " [12. 13. 14. 15.]\n",
            " [16. 17. 18. 19.]\n",
            " [20. 21. 22. 23.]]\n",
            "<NDArray 6x4 @cpu(0)>\n",
            "load into [gpu(0)]\n",
            "output: [\n",
            "[[ 0.  1.  2.  3.]\n",
            " [ 4.  5.  6.  7.]\n",
            " [ 8.  9. 10. 11.]\n",
            " [12. 13. 14. 15.]\n",
            " [16. 17. 18. 19.]\n",
            " [20. 21. 22. 23.]]\n",
            "<NDArray 6x4 @gpu(0)>]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9YUj9xS6ivvQ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def train_batch(X, y, gpu_params, ctx, lr):\n",
        "    gpu_Xs, gpu_ys = split_and_load(X, ctx), split_and_load(y, ctx)\n",
        "    with autograd.record():\n",
        "        ls = [loss(lenet(gpu_X, gpu_W), gpu_y)\n",
        "        for gpu_X, gpu_y, gpu_W in zip(gpu_Xs, gpu_ys, gpu_params)]\n",
        "    for l in ls:\n",
        "        l.backward()\n",
        "    for i in range(len(gpu_params[0])):\n",
        "        allreduce([gpu_params[c][i].grad for c in range(len(ctx))])\n",
        "    for param in gpu_params:\n",
        "        d2l.sgd(param, lr, X.shape[0])"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zgvXVYAfk3O_",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def train(num_gpus, batch_size, lr):\n",
        "    train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)\n",
        "    ctx = [mx.gpu(i) for i in range(num_gpus)]\n",
        "    print('running on:', ctx)\n",
        "    gpu_params = [get_params(params, c) for c in ctx]\n",
        "    for epoch in range(4):\n",
        "        start = time.time()\n",
        "        for X, y in train_iter:\n",
        "            train_batch(X, y, gpu_params, ctx, lr)\n",
        "            nd.waitall()\n",
        "        train_time = time.time() - start \n",
        "\n",
        "        def net(x):\n",
        "            return lenet(x, gpu_params[0])\n",
        "        test_acc = d2l.evaluate_accuracy(test_iter, net, ctx[0])\n",
        "        print('epoch %d, time %.1f sec, test acc %.2f' % (epoch + 1, train_time, test_acc))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NkcKCSZJmKK4",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 105
        },
        "outputId": "f23e7a03-54c8-4c93-f759-5cdaf8236bb7"
      },
      "source": [
        "train(num_gpus=1, batch_size=256, lr=0.2)"
      ],
      "execution_count": 49,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "running on: [gpu(0)]\n",
            "epoch 1, time 8.8 sec, test acc 0.10\n",
            "epoch 2, time 8.6 sec, test acc 0.55\n",
            "epoch 3, time 8.7 sec, test acc 0.70\n",
            "epoch 4, time 8.8 sec, test acc 0.72\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_eqUDxcLmQqh",
        "colab_type": "text"
      },
      "source": [
        "## 8.5 多 GPU 计算的简洁实现"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Gd6lLmarpkuO",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import d2lzh as d2l \n",
        "import mxnet as mx \n",
        "from mxnet import autograd, gluon, init, nd \n",
        "from mxnet.gluon import loss as gloss, utils as gutils, nn \n",
        "import time "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DvaSO9ioqQnq",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def resnet18(num_classes):\n",
        "    def resnet_block(num_channels, num_residuals, first_block=False):\n",
        "        blk = nn.Sequential()\n",
        "        for i in range(num_residuals):\n",
        "            if i == 0 and not first_block:\n",
        "                blk.add(d2l.Residual(num_channels, use_1x1conv=True, strides=2))\n",
        "            else:\n",
        "                blk.add(d2l.Residual(num_channels))\n",
        "        return blk \n",
        "\n",
        "    net = nn.Sequential()\n",
        "    net.add(nn.Conv2D(64, kernel_size=3, strides=1, padding=1), \n",
        "        nn.BatchNorm(), nn.Activation('relu'))\n",
        "    net.add(resnet_block(64, 2, first_block=True), \n",
        "        resnet_block(128, 2), \n",
        "        resnet_block(256, 2), \n",
        "        resnet_block(512, 2))\n",
        "    net.add(nn.GlobalAvgPool2D(), nn.Dense(num_classes))\n",
        "    return net \n",
        "\n",
        "net = resnet18(10)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7tRBBcyAr_aK",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "ctx = [mx.gpu(0)]\n",
        "net.initialize(init=init.Normal(sigma=0.01), ctx=ctx)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "q0BWC3L6s91P",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 263
        },
        "outputId": "a43da49a-51fc-453d-e875-57cd784843f8"
      },
      "source": [
        "x = nd.random.uniform(shape=(4, 1, 28, 28))\n",
        "gpu_x = gutils.split_and_load(x, ctx)\n",
        "net(gpu_x[0])"
      ],
      "execution_count": 53,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "\n",
              "[[-5.7465440e-06 -3.3317165e-06 -1.3032401e-06 -1.3137640e-06\n",
              "  -3.0325145e-06 -5.3874373e-06 -4.6628379e-06  1.3724061e-06\n",
              "  -2.9884850e-06  2.7169640e-06]\n",
              " [-6.0186107e-06 -3.9834722e-06 -1.0086044e-06 -1.3575535e-06\n",
              "  -2.4883661e-06 -5.1102088e-06 -5.1461625e-06  1.0052943e-06\n",
              "  -2.8755123e-06  2.1098490e-06]\n",
              " [-5.5376263e-06 -3.8441981e-06 -1.0446329e-06 -8.3044563e-07\n",
              "  -3.0834099e-06 -4.9034970e-06 -4.6071827e-06  6.2406548e-07\n",
              "  -2.4988947e-06  2.3911816e-06]\n",
              " [-6.0026546e-06 -3.7010204e-06 -9.8344708e-07 -9.2437887e-07\n",
              "  -2.3985406e-06 -5.5074138e-06 -4.9411633e-06  1.3234688e-06\n",
              "  -3.3002682e-06  3.2781813e-06]]\n",
              "<NDArray 4x10 @gpu(0)>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 53
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "l8MyxlwmtfpL",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 105
        },
        "outputId": "9127e9c9-6cd7-4c77-85d3-f7f4f05c5153"
      },
      "source": [
        "weight = net[0].params.get('weight')\n",
        "\n",
        "try:\n",
        "    weight.data()\n",
        "except RuntimeError:\n",
        "    print('not initialized on', mx.cpu())\n",
        "weight.data(ctx[0])[0]"
      ],
      "execution_count": 54,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "\n",
              "[[[-0.00481301  0.01521761 -0.00545954]\n",
              "  [ 0.01197755 -0.00073117  0.00322847]\n",
              "  [-0.00564734 -0.00043122  0.00359513]]]\n",
              "<NDArray 1x3x3 @gpu(0)>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 54
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "E-6R45rYuoSz",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def train(num_gpus, batch_size, lr):\n",
        "    train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)\n",
        "    ctx = [mx.gpu(i) for i in range(num_gpus)]\n",
        "    print('running on:', ctx)\n",
        "    net.initialize(init=init.Normal(sigma=0.01), ctx=ctx, force_reinit=True)\n",
        "    trainer = gluon.Trainer(\n",
        "        net.collect_params(), 'sgd', {'learning_rate': lr})\n",
        "    loss = gloss.SoftmaxCrossEntropyLoss()\n",
        "    for epoch in range(4):\n",
        "        start = time.time()\n",
        "        for X, y in train_iter:\n",
        "            gpu_Xs = gutils.split_and_load(X, ctx)\n",
        "            gpu_ys = gutils.split_and_load(y, ctx)\n",
        "            with autograd.record():\n",
        "                ls = [loss(net(gpu_X), gpu_y)\n",
        "                      for gpu_X, gpu_y in zip(gpu_Xs, gpu_ys)]\n",
        "            for l in ls:\n",
        "                l.backward()\n",
        "            trainer.step(batch_size)\n",
        "        nd.waitall()\n",
        "        train_time = time.time() - start\n",
        "        test_acc = d2l.evaluate_accuracy(test_iter, net, ctx[0])\n",
        "        print('epoch %d, time %.1f sec, test acc %.2f' % (\n",
        "            epoch + 1, train_time, test_acc))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "55607tXgviAF",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 105
        },
        "outputId": "f0105fff-6415-462d-84a2-06e84841fa0e"
      },
      "source": [
        "train(num_gpus=1, batch_size=256, lr=0.1)"
      ],
      "execution_count": 56,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "running on: [gpu(0)]\n",
            "epoch 1, time 91.7 sec, test acc 0.86\n",
            "epoch 2, time 84.9 sec, test acc 0.91\n",
            "epoch 3, time 84.7 sec, test acc 0.90\n",
            "epoch 4, time 84.9 sec, test acc 0.93\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0EIouY7FwzrY",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def train(num_gpus, batch_size, lr):\n",
        "    train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)\n",
        "    ctx = [mx.gpu(i) for i in range(num_gpus)]\n",
        "    print('running on:', ctx)\n",
        "    net.initialize(init=init.Normal(sigma=0.01), ctx=ctx, force_reinit=True)\n",
        "    trainer = gluon.Trainer(net.collect_params(), 'sgd', {'learning_rate': lr})\n",
        "    loss = gloss.SoftmaxCrossEntropyLoss()\n",
        "    for epoch in range(4):\n",
        "        start = time.time()\n",
        "        for X, y in train_iter:\n",
        "            gpu_Xs = gutils.split_and_load(X, ctx)\n",
        "            gpu_ys = gutils.split_and_load(y, ctx)\n",
        "            with autograd.record():\n",
        "                ls = [loss(net(gpu_X), gpu_y) for gpu_X, gpu_y in zip(gpu_Xs, gpu_ys)]\n",
        "            for l in ls:\n",
        "                l.backward()\n",
        "            trainer.step(batch_size)\n",
        "        nd.waitall()\n",
        "        train_time = time.time() - start \n",
        "        test_acc = d2l.evaluate_accuracy(test_iter, net, ctx[0])\n",
        "        print('epoch %d, time %.1f sec, test acc %.2f' %(epoch + 1, train_time, test_acc))\n",
        "                "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5kCbDP34zHOE",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 105
        },
        "outputId": "167f0bc7-075c-4128-d3a9-68cb9cd22269"
      },
      "source": [
        "train(num_gpus=1, batch_size=256, lr=0.1)"
      ],
      "execution_count": 60,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "running on: [gpu(0)]\n",
            "epoch 1, time 85.9 sec, test acc 0.89\n",
            "epoch 2, time 84.7 sec, test acc 0.92\n",
            "epoch 3, time 84.5 sec, test acc 0.93\n",
            "epoch 4, time 85.0 sec, test acc 0.91\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}